/** 
* @file allocator.cpp
* Implementation of kernel memory allocation engine (kernel heap).
* Inspired by JamesM's kernel development tutorials.
*
* Copyrights 2010 Michal Saman, m.saman@designplus.cz.
* This source code is release under the Apache License 2.0.
* http://www.apache.org/licenses/LICENSE-2.0
*/

#include <env.h>
#include <mem/allocator.h>
#include <mem/vmu.h>
#include <hal/hal.h>

/**
 * Init heap allocator engine.
 *
 * @param	heap_size		Size in bytes we want to allocate for heap.
 * @param	heap_start		Memory addres for start of heap
 *
 * @returns	Pointer to heap descriptor (structure).
 */
heap_t* init_allocator(uint32_t heap_size, uint32_t heap_start) {
	uint32_t i, data_heap_start, end_address;
	heap_t *heap;

	end_address = heap_start + heap_size;

	if (end_address & 0xFFFFF000 != 0) {
		end_address &= 0xFFFFF000;
		end_address += 0x1000;
	}

	//disable_interrupts();
	vmu_do_paging(false);

	for (i = heap_start; i <= end_address; i += PAGE_SIZE) {
		if (vmu_alloc_page( vmu_get_page(i) ) != true) {
			// we don't have space to allocate whole heap
			heap_size = (uint32_t)(i - heap_start);
			if (heap_size > 0)
				kprintf("Warning: lack of memory to create heap of required size,\nheap reduced to: %u kB.\n", heap_size/1024);
			break;
		}
	}

	vmu_do_paging(true);

	if (heap_size > sizeof(heap_t)) {
		heap = (heap_t*) heap_start;
		data_heap_start = heap_start + sizeof(heap_t);
		create_heap(heap, data_heap_start, end_address, 0, 0, 0);
	} else {
		kprintf("Critical error: lack of memory to create heap!\n");
		return NULL;
	}

	//enable_interrupts();

	return heap;
}

/**
* Lessthan_t implementation for compare size of block when inserting to heap directory ordered array.
* @see lessthan_t
*/
static int8_t header_t_lessthan(void* a, void* b) {
	//kprintf("Checking size %u < %u\n", ((header_t*)a)->size, ((header_t*)b)->size );
	return (((header_t*)a)->size < ((header_t*)b)->size) ? 1:0;
}

/**
* Create memory heap.
* @param start Start address where to place memory heap.
* @param end End address of memory heap.
* @param max Maximum end address when expanding heap.
* @param supervisor Supervisor mode flag.
* @param readonlu Readonlu mode flag.
*/
void create_heap(heap_t* heap, uint32_t start, uint32_t end, uint32_t max, uint8_t supervisor, uint8_t readonly) {
	// Init heap directory.
	heap->directory = ord_array_create((void*)start, HEAP_INDEX_SIZE, &header_t_lessthan);
	// Calculate sizeof heap directory (ordered array) to determine address where we can start putting data.
	start += sizeof(type_t)*HEAP_INDEX_SIZE;

	// Make sure the start address is page-aligned.
	if (start & 0xFFFFF000 != 0) {
		start &= 0xFFFFF000;
		start += 0x1000;
	}

	// write data start, data end, max address into heap structure
	heap->start_address = start;
	heap->end_address = end;
	heap->max_address = max;
	heap->size = (uint32_t)(end - start);
	heap->supervisor = supervisor;
	heap->readonly = readonly;

	// set out heap to one large free space (one deep black hole :)
	header_t* one_hole = (header_t*) start;
	one_hole->size = (uint32_t)(end - start);
	one_hole->magic = HEAP_MAGIC;
	one_hole->is_hole = 1;
	ord_array_insert(one_hole, &heap->directory);			// insert pointer to hole header to heap index
}

/**
* Extend size of heap.
* @param new_size New size of heap.
* @param heap Heap to externd.
*/
static bool expand(uint32_t new_size, heap_t *heap) {
	if (new_size <= heap->size)
		return false;	// sanity check

	if (new_size&0xFFFFF000 != 0) {
		new_size &= 0xFFFFF000;
		new_size += 0x1000;
	}

	if (heap->start_address + new_size > heap->max_address)
		return false;		// check if not larger than max_address

	uint32_t i = heap->size;
	disable_interrupts();
	vmu_do_paging(false);
	while (i<new_size) {
		if (vmu_alloc_page( vmu_get_page((virtual_addr)heap->start_address+i))) {
			i+=PAGE_SIZE;
		} else {
			break;  // out of space
		}
	}
	vmu_do_paging(true);
	enable_interrupts();
	if (i>heap->size) {
		heap->end_address = heap->start_address + i;
		heap->size = (uint32_t)(heap->end_address - heap->start_address);
		return true;
	} else {
		return false;
	}

}

/**
* Functions iterates withing heap directory and try to find the first smallest hole.
* @param size Size of hole to find in bytes.
* @param page_align If set to 1 requested the memory be page-aligned.\
* @param heap Heap whete to search.
*/
static int32_t find_smallest_hole(uint32_t size, uint8_t page_align, heap_t *heap) {
	uint32_t i = 0;

	//kprintf("Dir size: %u\n", heap->directory.size);
	while (i < heap->directory.size) {
		header_t *header = (header_t*) ord_array_lookup(i, &heap->directory);
		//kprintf("Checking hole at: 0x%x, size: %u, for size: %u\n", header, header->size, size);
		// page-aligned request
		if (page_align > 0) {
			uint32_t loc = (uint32_t) header;
			int32_t offset = 0;
			if ( ((loc+sizeof(header_t)) & 0xFFFFF000) != 0) {
				offset = PAGE_SIZE - (loc + sizeof(header_t))%PAGE_SIZE;
			}
			// fit hole size to fit page
			int32_t hole_size = (int32_t)(header->size - offset);
			//kprintf("offset: 0x%x, loc: 0x%x\n", offset, (loc+sizeof(header_t)) & 0xFFFFF000 );
			if ((hole_size >= (uint32_t)size) && (header->is_hole)) {
				//kprintf("I: %i, hole_size: %i -- ", i, hole_size);
				break;
			}
		} else if ((header->size >= size) && header->is_hole) {
			//kprintf("FOUND!!\n");
			break;
		}
		i++;
	}

	//kprintf("heap size: %i\n", heap->directory.size);

	// if we go to end, we didnt find free hole - out of mem...
	if ((i == heap->directory.size) && (heap->directory.size>0))
		return 	-1;
	else
		return 	 i;  			// return position index of free hole inside heap directory
}

/**
* Allocate block of memory in declared heap.
* @param size Size of memory block to allocate (in bytes).
* @param page_align Should be base address aligned to page?
* @param heap Heap where allocate memory.
*/
void *alloc(uint32_t size, uint8_t page_align, heap_t *heap) {
	kprintf("Heap size: %u, heap: 0x%x\n", heap->directory.size, heap);

	// add header and footer to allocation size
	uint32_t alloc_size = size + sizeof(header_t) + sizeof(footer_t);

	// find position of hole in heap directory
	int32_t i = find_smallest_hole(alloc_size, page_align, heap);
	//kprintf("Hole index: %u, %u\n", i, page_align);
	if (i==-1) { 				// we dont find suitable free hole
		kprintf("HEAP OUT OF SPACE!!\n");
		return NULL;
		/**
		uint32_t old_length = (uint32_t)(heap->end_address - heap->start_address);
		// save old end address before expanding
		uint32_t old_end_address = heap->end_address;
		// we need to expand heap size
		if (expand(old_length, heap) != false)
			return 0;			// HEAP TOTALLY OUT OF SPACE, SORRY !!
		uint32_t new_length = (uint32_t)(heap->end_address - heap->start_address);
		// try to find endmost header
		i = 0;
		int32_t idx = -1; uint32_t value = 0x0;
		while ((uint32_t)i<heap->directory.size) {
			type_t tmp = ord_array_lookup(i, &heap->directory);
			if ((uint32_t)tmp > value) {
				value = (uint32_t)tmp;
				idx = i;
			}
			i++;
		}
		// if we didnt find any headers, we neeed to add one.
		if (idx == -1) {
			header_t *header = (header_t *)old_end_address;
			header->magic = HEAP_MAGIC;
			header->size = (uint32_t)(new_length - old_length);
			header->is_hole = 1;
			footer_t *footer = (footer_t *)(uint32_t)(old_end_address + header->size - sizeof(footer_t));
			footer->header = header;
			ord_array_insert((void *)header, &heap->directory);
		} else {
			// last header needs adjusting
			header_t * header = (header_t*)ord_array_lookup(idx, &heap->directory);
			header->size += (uint32_t)(new_length - old_length);
			// rewrite footer
			footer_t *footer = (footer_t*)(header+header->size-sizeof(footer_t));
			footer->header = header;
			footer->magic = HEAP_MAGIC;
		}
		// heap is expanded, recurse and call the function again.
		return alloc(size, page_align, heap);
		*/
	}

	// save header info of finded hole
	header_t *orig_hole_header = (header_t*) ord_array_lookup((uint32_t)i, &heap->directory);
	if (!orig_hole_header) {
		kprintf("Don't found index in array, exiting! addr: 0x%x\n", orig_hole_header);
		return NULL; 			// we don't found index in array
	}

	uint32_t orig_hole_pos = (uint32_t)orig_hole_header;	// save addr of header
	uint32_t orig_hole_size = orig_hole_header->size;		// save hole size
	if ((uint32_t)(orig_hole_size-alloc_size) <= (sizeof(header_t)+sizeof(footer_t))) {					// determine if there will be some usefull space after allocation hole (at least 1 byte)
		size += (uint32_t)(orig_hole_size-alloc_size);													// there will be no usefull space after alloc, set space to alloc to whole hole size
		alloc_size = orig_hole_size;
	}

	// determine if we need page-align data
	// set new hole before our block to allocation
	if (page_align && (orig_hole_size & 0xFFFFF000)) {
		//kprintf("allocator, page aligned hole request\n");
		uint32_t new_location = (uint32_t)(orig_hole_pos + PAGE_SIZE - (orig_hole_pos & 0xFFF) - sizeof(header_t));
		//kprintf("Orig. hole loc: 0x%x, new hole loc: 0x%x\n", orig_hole_pos, new_location);
		header_t *hole_header = (header_t*) orig_hole_pos;
		hole_header->size = (uint32_t)(PAGE_SIZE - (orig_hole_pos & 0xFFF) - sizeof(header_t));
		hole_header->magic = HEAP_MAGIC;
		hole_header->is_hole = 1;
		footer_t * hole_footer = (footer_t *)((uint32_t)new_location - sizeof(footer_t));
		hole_footer->magic = HEAP_MAGIC;
		hole_footer->header = hole_header;
		// set new position and size to our aligned hole
		orig_hole_pos = new_location;
		orig_hole_size = (uint32_t)(orig_hole_size - hole_header->size);
	} else {
		// Else we dont need this orig hole anymore, delete it
		ord_array_delete((uint32_t)i, &heap->directory);
	}

	// Write our allocation block (rewrite orig hole)
	header_t *block_header = (header_t*)orig_hole_pos;
	block_header->magic = HEAP_MAGIC;
	block_header->is_hole = 0;
	block_header->size = alloc_size;
	footer_t *block_footer = (footer_t*)((uint32_t)block_header+sizeof(header_t)+size);
	block_footer->magic = HEAP_MAGIC;
	block_footer->header = block_header;
	//kprintf("alloc size: %u (size: %u), ft: 0x%x, hd: 0x%x, adr: 0x%x\n", alloc_size, size, block_footer, block_header, block_header+sizeof(header_t));

	// If new hole have a positive size we neeed to write it after allocated block
	if ((int32_t)(orig_hole_size - alloc_size) > 0) {
		header_t *hole_header =(header_t*)(orig_hole_pos+sizeof(header_t)+size+sizeof(footer_t));
		hole_header->magic = HEAP_MAGIC;
		hole_header->is_hole = 1;
		hole_header->size = (uint32_t)(orig_hole_size - alloc_size);
		footer_t * hole_footer = (footer_t*)((uint32_t)hole_header+orig_hole_size-alloc_size-sizeof(footer_t));
		if ((uint32_t)hole_footer < heap->end_address) {
			hole_footer->magic = HEAP_MAGIC;
			hole_footer->header = hole_header;
		}
		// insert new hole to heap direcotry
		ord_array_insert((void*)hole_header, &heap->directory);
		kprintf("Heap size2: %u, heap: 0x%x\n", heap->directory.size, heap);
	}
	// return poiner to allocated space
	return (void*) ((uint32_t)block_header+sizeof(header_t));
}

/**
* Free block of memory in declared heap.
* @param p 		Pointer to variable allocated previously by allocator.
* @param heap 	Heap where variable was allocated.
*/
void free(void *p, heap_t *heap) {
	if (p==0) 
		return;
	// get header and footer associated with the pointer
	header_t *header = (header_t*) ((uint32_t)p - sizeof(header_t));
	footer_t *footer = (footer_t*) ((uint32_t)p - sizeof(header_t) + header->size - sizeof(footer_t));
	//kprintf("--- Allocator: free() - dealoc size: %i, ft: 0x%x, hd: 0x%x, m: 0x%x\n", header->size, footer, header, footer->magic);
	// test if this is really heap allocated block
	if (header->magic!=HEAP_MAGIC || header->magic!=HEAP_MAGIC) {
		//kprintf("Allocator: free() - wrong magic, exiting!\n");
		return;
	}

	header->is_hole = 1;
	uint8_t to_add = 1;

	// UNIFY LEFT
	// check if before us is another footer, try to load left footer
	footer_t *l_footer = (footer_t*)((uint32_t)header - sizeof(footer_t));
	//kprintf("l_footer magic: 0x%x\n", l_footer->magic);
	if (l_footer->magic == HEAP_MAGIC && l_footer->header->is_hole == 1) {
		uint32_t old_size = header->size;
		//kprintf("Allocator: free() - leftU, my-size: %u, left-hole size: %u, dir size: %u \n", old_size, l_footer->header->size, heap->directory.size);
		header = l_footer->header;			// rewrite our header with the new one.
		footer->header = header;			// set footer-header to pointers new header.
		header->size += old_size;			// expand size
#ifdef ALLOC_FREE_REORDER
		to_add = 0;
		uint32_t i = 0;
		while ((i < heap->directory.size) &&
			(ord_array_lookup(i, &heap->directory) != (void*) header))
			i++;
		ord_array_delete(i, &heap->directory);
		ord_array_insert((void*)header, &heap->directory);
#else
		to_add = 0;
#endif
	}

	// UNIFY RIGHT
	// check if behind us is header with another hole
	header_t *r_header = (header_t*)((uint32_t)footer + sizeof(footer_t));
	if (r_header->magic==HEAP_MAGIC && r_header->is_hole==1) {
		//kprintf("Allocator: free() - rightU, my-size: %u, right-hole size: %u, dir size: %u \n", header->size, r_header->size, heap->directory.size);
		header->size += r_header->size;   // increase our size
		// set new footer location
		footer = (footer_t*)((uint32_t)r_header + r_header->size - sizeof(footer_t));
		footer->header = header;
		footer->magic = HEAP_MAGIC;
		// find and remove this header from heap directory
		uint32_t i = 0;
		while ((i < heap->directory.size) && 
			(ord_array_lookup(i, &heap->directory) != (void*)r_header)) 
			i++;
		//kprintf("rightU: index to delete: %i\n", i);
		// check if we find header
		if (i>=heap->directory.size) {
			kprintf("FUCK\n");
			return;
		}
		// remove it
		ord_array_delete(i, &heap->directory);
	}

	// if the footer is at end of heap, we can contract size of heap.
	// TODO

	if (to_add == 1)
		ord_array_insert((void*)header, &heap->directory);
}

/**
 * Show debug information about allocator engine.
 * @param heap Heap about to show inormation.
 */
void allocator_info(heap_t *heap) {
	kprintf("Allocator directory at: 0x%x\n", heap->directory);
	kprintf("Allocator directory entries count: %u\n", heap->directory.size);
	kprintf("Allocator directory max entries: %u\n", heap->directory.max_size);
	uint32_t ord_array_size = heap->directory.max_size * sizeof(type_t);
	kprintf("Allocator directory structure size: %u (0x%x) bytes\n", ord_array_size, ord_array_size);
	kprintf("Allocator data start at: 0x%x\n", heap->start_address);
	uint32_t i, free = 0;
	for (i=0; i < heap->directory.size; i++) {
		kprintf("Hole #%i at: 0x%x, size: %u b, data size: %u b\n", i+1, heap->directory.entries[i], ((header_t*)heap->directory.entries[i])->size, ((header_t*)heap->directory.entries[i])->size - sizeof(header_t) - sizeof(footer_t) );
		free = free + ((header_t*)heap->directory.entries[i])->size;
	}
	kprintf("Total size: ~%u MB (%u b), start at: 0x%x, end at: 0x%x\n", heap->size/1024/1024, heap->size, heap->start_address, heap->end_address );
	kprintf("Free space: ~%u MB (%u b)\n", free/1024/1024, free);
}
